package com.lansetech.framework.redis.peak.shaving;

import org.springframework.data.redis.connection.jedis.JedisConnectionFactory;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

public class PeakShavingRunnableContext {

    private final JedisConnectionFactory factory;

    private final AtomicReference<RunnableState> state;

    private final Map<Integer, Lock> runnerLockMap = new HashMap<>();
    private final Map<Integer, Condition> runnerConditionMap = new HashMap<>();

    private final BlockingQueue<PeakShavingRunnable> awaitBlockingQueue = new LinkedBlockingQueue<>();

    private final PeakShavingConfig config;

    private CountDownLatch shutdownLatch;

    public PeakShavingRunnableContext(JedisConnectionFactory factory, PeakShavingConfig config) {
        this.factory = factory;
        this.state = new AtomicReference<>();
        this.state.set(RunnableState.INIT);
        this.config = config;
    }

    public JedisConnectionFactory getConnectionFactory() {
        return factory;
    }

    public PeakShavingConfig getConfig(){
        return config;
    }

    public RunnableState getState() {
        return state.get();
    }

    public void setState(RunnableState state) {
        this.state.set(state);
    }

    public void registerConsumer(int index) {
        Lock lock = new ReentrantLock();
        this.runnerLockMap.put(index, lock);
        this.runnerConditionMap.put(index, lock.newCondition());
    }

    public CountDownLatch getShutdownLatch() {
        return shutdownLatch;
    }

    public void setShutdownLatch(CountDownLatch shutdownLatch) {
        this.shutdownLatch = shutdownLatch;
    }

    public Condition obtainLockCondition(int index) {
        return this.runnerConditionMap.get(index);
    }

    public Lock obtainLock(int index) {
        return this.runnerLockMap.get(index);
    }

    public void wakeMeUp(PeakShavingRunnable runner) {
        awaitBlockingQueue.add(runner);
    }

    public BlockingQueue<PeakShavingRunnable> getAwaitBlockingQueue() {
        return awaitBlockingQueue;
    }
}
